# -*- coding: utf-8 -*-
# Created Time : 2022/1/5 15:19
# Author:Zhou
import os

import h5py
import matplotlib.pyplot as plt
import numpy as np
from ANC_tools import sin_wave, awgn, NPP_1,NPP_2,NPP_3,NPP_4,NPP_5, logisticChaotic

# some settings
# 生成训练数据
rms = 1.0
n_tr_ex = 256

mix = sin_wave(A=1, f=100, fs=16000, phi=0, t=1) \
      + sin_wave(A=1, f=100, fs=16000, phi=90, t=1) \
      + sin_wave(A=0.0, f=200, fs=16000, phi=180, t=1)

# mix = logisticChaotic(0.9, 4, 48000)
# mix = awgn(mix, 40)
# sph = NPP_1(mix)
sph = sin_wave(A=1, f=100, fs=16000, phi=0, t=1)

# normalize
c = rms * np.sqrt(mix.size / np.sum(mix ** 2))
mix *= c
sph *= c
# plt.figure()
# plt.subplot(2, 1, 1)
# plt.plot(mix)
# plt.subplot(2, 1, 2)
# plt.plot(sph)
# plt.show()




for idx in range(n_tr_ex):  # n_tr_ex is the number of training examples
    # generate a noisy mixture
    # mix stores a noisy mixture utterance
    # sph the corresponding clean speech utterance.
    # 50-500之间的随机频率
    # speech enhancement 的任务是输入mix(speech+noise)，得到接近sph的声音
    # ANC的任务是  输入reference (x(n)，单频参考噪声), 得到接近初级噪声(P(n))的信号

    # print('generating: sample_{}'.format(idx))
    # save file
    filepath = 'datasets/tr/'
    filename = 'tr_{}.ex'.format(idx)
    writer = h5py.File(os.path.join(filepath, filename), 'w')
    writer.create_dataset('mix', data=mix.astype(np.float32), shape=mix.shape, chunks=True)
    writer.create_dataset('sph', data=sph.astype(np.float32), shape=sph.shape, chunks=True)
    writer.close()

for i in range(n_tr_ex):
    print("../data/datasets/tr/tr_{}.ex".format(i))
